/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022.
 * Description: kexec zImage support
 * Author: songzhezhe
 * Create: 2022-03-07
 */

#define pr_fmt(fmt)	"kexec_file(zImage): " fmt

#include <linux/err.h>
#include <linux/errno.h>
#include <linux/kernel.h>
#include <linux/kexec.h>
#include <linux/pe.h>
#include <linux/string.h>
#include <linux/verification.h>
#include <asm/byteorder.h>
#ifdef CONFIG_RTOS_HAL_ARM32_SUPPORT_KEXEC_FILE
#include <asm/sections.h>
#else
#include <asm/cpufeature.h>
#include <asm/image.h>
#endif
#include <asm/memory.h>

#define MAX(a, b) ((a) > (b) ? (a) : (b))

static int zimage_probe(const char *kernel_buf, unsigned long kernel_len)
{
	/* check nothing for zimage */
	return 0;
}

/* assume the maximum kernel compression ratio is 10, and just to be safe */
#define MAX_COMPRESSION_RATIO  10
/* text_offset: The image load offset in LSB byte order. */
#ifdef CONFIG_RTOS_HAL_ARM32_SUPPORT_KEXEC_FILE
#define ARM_TEXT_OFFSET  0x8000
#endif
#define ARM64_TEXT_OFFSET  0x80000

static unsigned long get_image_size(unsigned long val)
{
	unsigned long len = val;
	unsigned long kexec_image_size;
	unsigned long estimate_size;

	/*
	 * The zImage length does not include its stack (4k) or its
	 * malloc space (64k).  Include this.
	 */
	len += 0x11000;
	pr_debug("zImage requires 0x%08lx bytes\n", len);
	/*
	 * assume the maximum kernel compression ratio is 10.
	 * some open source just simply use SZ_32M, like efi_stub.
	 * to keep safe, we select the max one.
	 */
	estimate_size = len * MAX_COMPRESSION_RATIO;
	kexec_image_size =  MAX(estimate_size, SZ_32M);
	pr_debug("kernel image requires 0x%08lx bytes\n", kexec_image_size);

	return kexec_image_size;
}

static void *zimage_load(struct kimage *image,
				char *kernel, unsigned long kernel_len,
				char *initrd, unsigned long initrd_len,
				char *cmdline, unsigned long cmdline_len)
{
	struct kexec_buf kbuf;
	unsigned long text_offset;
	struct kexec_segment *kernel_segment = NULL;
	int ret;

	/* Load the kernel */
	kbuf.image = image;
	kbuf.buf_min = 0;
#ifdef CONFIG_RTOS_HAL_ARM32_SUPPORT_KEXEC_FILE
	kbuf.buf_max = virt_to_phys(high_memory);
#else
	kbuf.buf_max = ULONG_MAX;
#endif
	kbuf.top_down = false;

	kbuf.buffer = kernel;
	kbuf.bufsz = kernel_len;
	kbuf.mem = KEXEC_BUF_MEM_UNKNOWN;
	kbuf.memsz = ALIGN(get_image_size(kernel_len), PAGE_SIZE);
#ifdef CONFIG_RTOS_HAL_ARM32_SUPPORT_KEXEC_FILE
	text_offset = ARM_TEXT_OFFSET;
	kbuf.buf_align = SZ_2M;
#else
	text_offset = ARM64_TEXT_OFFSET;
	kbuf.buf_align = MIN_KIMG_ALIGN;
#endif
	/* Adjust kernel segment with TEXT_OFFSET */
	kbuf.memsz += text_offset;

	ret = kexec_add_buffer(&kbuf);
	if (ret)
		return ERR_PTR(ret);

	kernel_segment = &image->segment[image->nr_segments - 1];
#ifdef CONFIG_RTOS_HAL_ARM32_SUPPORT_KEXEC_FILE
	kernel_segment->mem = virt_to_phys(_text) - kaslr_offset();
#else
	kernel_segment->mem += text_offset;
	kernel_segment->memsz -= text_offset;
#endif
	image->start = kernel_segment->mem;

#ifdef CONFIG_RTOS_HAL_ARM32_SUPPORT_KEXEC_FILE
	pr_debug("Loaded zImage kernel at 0x%pa bufsz=0x%lx memsz=0x%x\n",
#else
	pr_debug("Loaded zImage kernel at 0x%pa bufsz=0x%lx memsz=0x%lx\n",
#endif

				&kernel_segment->mem, kbuf.bufsz,
				kernel_segment->memsz);

	/* Load additional data */
	ret = load_other_segments(image,
				kernel_segment->mem, kernel_segment->memsz,
				initrd, initrd_len, cmdline);

	return ERR_PTR(ret);
}

const struct kexec_file_ops kexec_zimage_ops = {
	.probe = zimage_probe,
	.load = zimage_load,
};
